# Copyright (c) 2024 Huawei Technologies Co., Ltd.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.


import platform
import importlib
import importlib.metadata

from .. import __version__ as version
from .. import (
    is_ms_available,
    is_torch_available,
    is_torch_npu_available,
)


def _grab_package_version(pkg_name: str) -> str:
    try:
        return importlib.metadata.version(pkg_name)
    except importlib.metadata.PackageNotFoundError:
        return "not installed"


def _format_dict(d):
    return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"


def run_env() -> None:
    """Get current environment information."""

    pt_version = "not installed"
    pt_npu_version = "not installed"
    pt_npu_available = False
    if is_torch_available():
        import torch

        pt_version = torch.__version__
        pt_npu_available = is_torch_npu_available()
        if pt_npu_available:
            import torch_npu

            pt_npu_version = torch_npu.__version__
    ms_version = "not installed"
    if is_ms_available():
        import mindspore

        ms_version = mindspore.__version__
    try:
        import openmind_hub

        om_hub_version = openmind_hub.__version__
    except Exception:
        om_hub_version = "not installed"
    info = {
        "`openmind` version": version,
        "`openmind_hub` version": om_hub_version,
        "Platform": platform.platform(),
        "Python version": platform.python_version(),
        "PyTorch version (NPU?)": f"{pt_version} ({pt_npu_version if pt_npu_available else pt_npu_available})",
        "Mindspore version": f"{ms_version}",
    }

    info["Mindformers version"] = _grab_package_version("mindformers")
    info["Transformers version"] = _grab_package_version("transformers")
    info["Accelerate version"] = _grab_package_version("accelerate")
    info["Datasets version"] = _grab_package_version("datasets")
    info["Evaluate version"] = _grab_package_version("evaluate")
    info["DeepSpeed version"] = _grab_package_version("deepspeed")
    info["Lm-evaluation-harness version"] = _grab_package_version("lm_eval")

    print(_format_dict(info))
